import os
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.models import Model
from tensorflow import keras
import module

""" load all of cifar """
x_train,y_train,_,_,label_names = module.load_data()
input_shape = x_train.shape[1:]
#劣化版LeNet-5

# model = keras.models.Sequential([
#     # 卷积层1
#     keras.layers.Conv2D(6, 5, input_shape=input_shape), # 使用6个5*5卷积核     
#     keras.layers.ReLU(),                # ReLU激活函数
#     keras.layers.MaxPooling2D(pool_size=2, strides=2),   
#     # 卷积层2
#     keras.layers.Conv2D(16, 5),         # 使用16个5*5的卷积
#     keras.layers.ReLU(),                # ReLU激活函数
#     keras.layers.MaxPooling2D(pool_size=2, strides=2),    
#     # 卷积层3
#     keras.layers.Conv2D(120, 5),        # 使用120个5*5的卷积     
#     keras.layers.ReLU(),                # ReLU激活函数
#     keras.layers.Flatten(),             # 2维变1维, 展平
#     # 全连接层1
#     keras.layers.Dense(84, activation='relu'),    # 120*84
#     # 全连接层2
#     keras.layers.Dense(len(label_names), activation='softmax')  # 84*1
# ])

# 加载VGG模型 keras.applications.vgg16，weights为imagenet表示使用imagenet预训练的模型，classes表示1000种分类
base_model = VGG16(include_top=False, weights='imagenet',input_shape = input_shape )
with tf.name_scope('output'):
    x = base_model.output
    x = keras.layers.Flatten()(x)
    x = keras.layers.Dense(84, activation='relu')(x)
    predictions  = keras.layers.Dense(len(label_names), activation='softmax')(x)
model = Model(inputs = base_model.inputs,outputs = predictions)

# model = keras.models.Sequential([
#     base_model,
#     keras.layers.Flatten(),
#     keras.layers.Dense(84, activation='relu'),
#     keras.layers.Dense(len(label_names), activation='softmax')
# ])
model.summary()

model.compile(optimizer=tf.keras.optimizers.SGD(lr=0.01),loss=tf.keras.losses.categorical_crossentropy, metrics=['accuracy',module.recall,module.precision])
validation_split=0.2 #表示在训练集中采用 20%作为验证集，
validation_freq=200 # 20个epoch输出一次验证集的学习效果
batch_size = 200
training_epochs = 40
log_dir="./logs" 
if not os.path.exists(log_dir):
    os.mkdir(log_dir)    # 创建保存目录
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1,profile_batch = 100000000)
history = model.fit(x_train, y_train,batch_size=batch_size,epochs=training_epochs,validation_split=validation_split, validation_freq=validation_freq,callbacks=[tensorboard_callback])
# plt.plot(history.history['loss'])    # 损失下降曲线
model.save(module.type+'_cs_model.h5')
# plt.show()
